'''
Created on Apr 3, 2012

@author: tel
'''
#site.addsitedir(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import site, os, shutil, threading, copy, time, numpy, sys, random, string
sys.path.append('/home/tel/workspace/pmcce/src')
#site.addsitedir('/home/tel/workspace/pmcce/src/')
#global vars: META, DB

from sqlalchemy import create_engine, MetaData, event, interfaces
from sqlalchemy.orm import sessionmaker, scoped_session, mapper, relationship

import config
import rules

file_names = ('run.prm', 'pK.out', 'energies.opp', 'head3.lst', 'head2.lst', 'head1.lst', 'fort.38', 'step2_out.pdb') 
rule_dict = {'run.prm':rules.RunRule, 
             'step2_out.pdb':rules.Step2Rule, 
             'head1.lst':rules.Head1Rule, 
             'head2.lst':rules.Head2Rule, 
             'head3.lst':rules.Head3Rule, 
             'energies.opp':rules.OppRule, 
             'fort.38':rules.FortRule, 
             'pK.out':rules.PKRule}

def ListSplit(l, s):
    splitlist = [l[i:i+len(l)/s] for i in xrange(0, len(l), len(l)/s)]
    if len(splitlist)>s:
        tmp = splitlist.pop(-1)
        splitlist[-1] = splitlist-1+tmp
    return splitlist

class FileList(object):
    def __init__(self, file_paths):
        self.file_paths = file_paths
        self.files = []
    
    def Open(self):
        for fi in self.file_paths:
            self.files.append(open(fi, 'r'))
    
    def Close(self):
        for fi in self.file_paths:
            fi.close()

def ListSearch(l1, l2):
    for item in l1:
        if item in l2:
            return item
    return False

def AbsToRel(root, fullpath):
    return fullpath[len(root)+1:]

def RelToAbs(root, relpath):
    return os.path.join(root, relpath)

def FindInBranch(root, branch, files):
    file_paths = []
    files = list(files)
    root = root.strip()
    root = root.rstrip('/')
    branch = branch.strip()
    branch = branch.rstrip('/')
    while len(branch) > len(root):
        dir_list = os.listdir(branch)
        flength = len(files)
        for i, fi in enumerate(reversed(files)):
            if fi in dir_list:
                file_paths.append(os.path.join(branch, fi))
                files.pop(flength-i-1)
        branch = os.path.split(branch)[0]
    return file_paths

def _fk_pragma_on_connect(dbapi_con, con_record):
    dbapi_con.execute('pragma foreign_keys = ON')
    dbapi_con.execute('pragma synchronous = OFF')

def MakeDB(meta, db, session, rootdir, protdirs, clevels, lock, event):
    hierarchy = clevels.hierarchy
    count = 0
    tcount = 0
    tennow = time.clock()
    tentimes = []
    for protdir in protdirs:
        hierarchy[0].Purge()    #equivalent to levels.Prot.Purge()
        now = time.clock()
        full = True
        for tup in (tup for tup in os.walk(os.path.join(rootdir,protdir)) if 'pK.out' in tup[2]):
            if '/full' in tup[0]:
                full = True
        if full==True:
            for tup in (tup for tup in os.walk(os.path.join(rootdir,protdir)) if 'pK.out' in tup[2]):
                data = {'Prot':[('pdb', protdir)]}
                
                try:
                    pdb_path = FindInBranch(rootdir, tup[0], ('%s.pdb' % protdir,))[0]
                except IndexError:
                    pdb_path = FindInBranch(rootdir, tup[0], ('prot.pdb',))[0]
                pdb_file = open(pdb_path, 'r')
                pdb_rule = rules.SeqadvRule(pdb_file)
                pdb_file.close()
                data['seqadv'] = pdb_rule.data
                
                file_paths = FindInBranch(rootdir, tup[0], file_names)
                file_list = FileList(file_paths)
                file_list.Open()
                for fi in file_list.files:
                    rule = rule_dict[fi.name.split('/')[-1]](fi)
                    fi.close()
                    data[rule.rtype] = rule.data
                lower_lvl = None
                for level in hierarchy:
                    level.Org(data)
                for level in hierarchy:
                    level.Table(meta=meta, db=db)
                for level in reversed(hierarchy):
                    level.SQLize()
                for level in reversed(hierarchy):
                    if lower_lvl != None:
                        level.AddChildren()
                        session.add(level.instances[0]) #bugfix for missing ATOMS at the start of some proteins
                    lower_lvl = level
                hierarchy[0].instances[0].Runs.append(hierarchy[1].instances[0])    #heirarchy[0] should be a prot copy, hierarchy[1] should be a run copy
                for level in hierarchy[1:]: #the [1:] slice leaves off the ProtData class so its instance persists over the Runs
                    level.Purge()
            if len(hierarchy[0].instances) > 0:
                looptime = time.clock() - now
                print "%d\t%s took:\t%s \t%s\t%s\t%s" % (tcount, clevels.ProtData.instances[0].pdb, looptime, id(hierarchy[0]), id(hierarchy[0].instances), id(hierarchy[0].instances[0]))
                tentimes.append(looptime)
                session.add(hierarchy[0].instances[0])
                count += 1
                tcount += 1
                if count % 3==0:
                    print "the last hundred non-sql loops took:\t%s" % numpy.sum(tentimes)
                    tentimes = []
                    event.set()
                    with lock:
                        session.commit()
                    print "the last hundred took:\t%s" % (time.clock()-tennow)
                    count = 0
                    tennow = time.clock()
    with lock:
        session.commit()

if __name__ == "__main__":
    
    # generate hash randomly with 6 letters or numbers 
    hash = ''.join(random.choice(string.ascii_uppercase + string.digits) for x in range(6))
    dbtype = 'mysq'
    indir = sys.argv[1]
    outdir = sys.argv[2]
    dbf = outdir.split('/')[-1]
    #dir = '/Users/tel/Documents/lab/test/db1/'
    #dir = '/home/tel/benchmark/dunbrack'
    #dir = '/home/tel/test/db1'
    #dbf = 'dunbrack.db'
    path = os.path.join(indir, dbf)
    tmppath = os.path.join('/tmp', hash+dbf)
    if dbtype == 'mysql':
        MySQLdb.server_init(('--basedir=/home/tel/test/mysql','--datadir=/home/tel/test/mysql/data', '--skip-grant-tables'))
        db = create_engine('mysql+mysqldb://tel:tel@localhost')
        db.execute("DROP DATABASE IF EXISTS foo") #create db)
        db.execute("CREATE DATABASE foo") #create db)
        db.execute("USE foo") # select new db
    else:
        try:
            os.remove(tmppath)
        except OSError:
            pass
        db = create_engine('sqlite:///' + tmppath)
#        db.execute('pragma main.page_size=4096')
#        db.execute('pragma main.lock_mode=EXCLUSIVE')
#        db.execute('pragma main.synchronous=NORMAL')
#        db.execute('pragma main.journal_mode=WAL')
#        db.execute('pragma main.cache_size=10000')
#        db.execute('pragma main.temp_store=MEMORY')
        event.listen(db, 'connect', _fk_pragma_on_connect)
    session = scoped_session(sessionmaker(bind=db))
    meta= MetaData(db)
    
    threads = []
    walker = os.walk(indir)
    rootwalk = walker.next()
    rootdir = rootwalk[0]
    rootsubdirs = rootwalk[1]
    numthreads = 1
    lock = threading.Lock()
    event = threading.Event()
    for protdirs in ListSplit(rootsubdirs,numthreads):
        import levels
        clevels = levels
        MakeDB(meta=meta,db=db,session=session,rootdir=rootdir,protdirs=protdirs,clevels=clevels,lock=lock,event=event)
#        t = threading.Thread(target=MakeDB, kwargs={'meta':meta,'db':db,'session':session,'rootdir':rootdir,'protdirs':protdirs,'clevels':clevels, 'lock':lock, 'event':event})
#        threads.append(t)
#        t.start()
#        event.wait()
        del sys.modules['levels']
#    while threading.enumerate()>0:
#        time.sleep(5)
    shutil.move(tmppath, outdir)